{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 用RNN做POS tagging\n",
    "\n",
    "姓名: \\[write-your-name-here\\]\n",
    "\n",
    "在这份作业中，你会用一个bidirectional recurrent neural network来做POS tagging。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import necessary libraries and set the random seeds\n",
    "\n",
    "import os\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F\n",
    "from torch.optim import lr_scheduler\n",
    "from torchtext import data\n",
    "import numpy as np\n",
    "import random\n",
    "from torch.utils.data import Dataset\n",
    "import time\n",
    "import shutil\n",
    "\n",
    "EMBEDDING_DIM = 300\n",
    "HIDDEN_DIM = 200\n",
    "\n",
    "\n",
    "USE_CUDA = torch.cuda.is_available()\n",
    "random.seed(53113)\n",
    "np.random.seed(53113)\n",
    "torch.manual_seed(53113)\n",
    "if USE_CUDA:\n",
    "    torch.cuda.manual_seed(53113)\n",
    "    \n",
    "BATCH_SIZE = 128\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "载入POS tagging训练和dev数据集。这些文件都是tab分隔的text和POS tag数据，"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_datasets():\n",
    "    text = data.Field(include_lengths=True)\n",
    "    tags = data.Field()\n",
    "#     train_data, val_data, test_data = data.TabularDataset.splits(path='Pytorch-POS-Tagger/RNN_Data_files/', train='train_data.tsv', validation='val_data.tsv', test='val_data.tsv', fields=[('text', text), ('tags', tags)], format='tsv')\n",
    "    train_data, val_data, test_data = data.TabularDataset.splits(path='./', train='train.txt', validation='dev.txt', test='dev.txt', fields=[('text', text), ('tags', tags)], format='tsv')\n",
    "\n",
    "    \n",
    "    batch_sizes = (BATCH_SIZE, BATCH_SIZE, BATCH_SIZE)\n",
    "    train_loader, dev_loader, test_loader = data.BucketIterator.splits((train_data, val_data, test_data), batch_sizes=batch_sizes, sort_key=lambda x: len(x.text))\n",
    "\n",
    "    text.build_vocab(train_data)\n",
    "    tags.build_vocab(train_data)\n",
    "    dataloaders = {'train': train_loader,\n",
    "                   'validation': dev_loader,\n",
    "                   'test': dev_loader}\n",
    "    return text, tags, dataloaders\n",
    "\n",
    "text, tags, dataloaders = load_datasets()\n",
    "text_vocab_size = len(text.vocab.stoi) + 1\n",
    "tag_vocab_size = len(tags.vocab.stoi) - 1   # = 42 (not including the <pad> token\n",
    "print(text_vocab_size)\n",
    "print(tag_vocab_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class POSDataset(Dataset):\n",
    "    def __init__(self, path, sen_vocab, tag_vocab):\n",
    "        super(POSDataset, self).__init__()\n",
    "        self.sen_vocab = sen_vocab\n",
    "        self.tag_vocab = tag_vocab\n",
    "        self.num_classes = tag_vocab.size()\n",
    "        sen_file = os.path.join(path, 'sentences.txt')\n",
    "        tag_file = os.path.join(path, 'tags.txt')\n",
    "        self.sentences = []\n",
    "        with open(sen_file, 'r') as f:\n",
    "            for line in f:\n",
    "                idxs = self.sen_vocab.toIdx(line.rstrip('\\n').split(' '))\n",
    "                tensor = torch.LongTensor(idxs)\n",
    "                self.sentences.append(tensor)\n",
    "\n",
    "        self.tags = []\n",
    "        with open(tag_file, 'r') as f:\n",
    "            for line in f:\n",
    "                idxs = self.tag_vocab.toIdx(line.rstrip('\\n').split(' '))\n",
    "                tensor = torch.LongTensor(idxs)\n",
    "                self.tags.append(tensor)\n",
    "\n",
    "        # making sure there are same number of sentences as tags.\n",
    "        assert(len(self.sentences) == len(self.tags))\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        sentence = self.sentences[index]\n",
    "        tags = self.tags[index]\n",
    "        return sentence, tags\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.sentences)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sequence_mask(sequence_length, max_len=None):\n",
    "    ''' Given a tensor of a sequence of lengths, create a mask of each length. \n",
    "    '''\n",
    "    if max_len is None:\n",
    "        max_len = sequence_length.data.max()\n",
    "    batch_size = sequence_length.size(0)\n",
    "    seq_range = torch.range(0, max_len - 1).long()\n",
    "    seq_range_expand = seq_range.unsqueeze(0).expand(batch_size, max_len)\n",
    "    if sequence_length.is_cuda:\n",
    "        seq_range_expand = seq_range_expand.cuda()\n",
    "    seq_length_expand = (sequence_length.unsqueeze(1)\n",
    "                         .expand_as(seq_range_expand))\n",
    "    return seq_range_expand < seq_length_expand\n",
    "\n",
    "\n",
    "# run one epoch of training\n",
    "def train(model, train_loader, loss_fn, optimizer, use_gpu=False):\n",
    "    model.train()  # Set model to training mode\n",
    "    running_loss = 0.0\n",
    "    running_corrects = 0\n",
    "    example_count = 0\n",
    "    step = 0\n",
    "    # Iterate over data.\n",
    "    for batch in train_loader:\n",
    "        sentences = batch.text[0].transpose(1, 0)\n",
    "        tags = batch.tags.transpose(1, 0)\n",
    "        ''' Implement the code to train the model. \n",
    "            - Prepare the input data (text, tags, mask) to the correct format and shape\n",
    "            - Run the forward method of the model\n",
    "            - Compute the loss\n",
    "            - Run backward on loss for back propagation\n",
    "            - Run the optimizer to update the model parameters. \n",
    "            - Compute the number of correct predictions\n",
    "        '''\n",
    "        # TODO\n",
    "        \n",
    "        \n",
    "        \n",
    "        step += 1\n",
    "        if step % 100 == 0:\n",
    "            print('loss: {}, running_corrects: {}, example_count: {}, acc: {}'.format(loss.item(), \n",
    "                            running_corrects, example_count, (running_corrects / example_count) * 100))\n",
    "        if step * batch_size >= 40000:\n",
    "            break\n",
    "    loss = running_loss / example_count\n",
    "    acc = (running_corrects / example_count) * 100\n",
    "    print(loss)\n",
    "    print(acc)\n",
    "    print('Train Loss: {:.4f} Acc: {:2.3f} ({}/{})'.format(loss, acc, running_corrects, example_count))\n",
    "    return loss, acc\n",
    "\n",
    "\n",
    "def validate(model, val_loader, loss_fn, use_gpu=False):\n",
    "    model.eval()  # Set model to evaluate mode\n",
    "    running_loss = 0.0\n",
    "    running_corrects = 0\n",
    "    example_count = 0\n",
    "    # Iterate over data.\n",
    "    with torch.no_grad():\n",
    "        for batch in val_loader:\n",
    "            sentences = batch.text[0].transpose(1, 0)\n",
    "            tags = batch.tags.transpose(1, 0) \n",
    "            ''' Similar to training, do the following to evaluate the model.  \n",
    "            - Prepare the input data (text, tags, mask) to the correct format and shape\n",
    "            - Run the forward method of the model\n",
    "            - Compute the loss\n",
    "            - Compute the number of correct predictions\n",
    "            '''\n",
    "            # TODO\n",
    "            \n",
    "\n",
    "    loss = running_loss / example_count\n",
    "    acc = (running_corrects / example_count) * 100\n",
    "    print(loss)\n",
    "    print(acc)\n",
    "    print('Validation Loss: {:.4f} Acc: {:2.3f} ({}/{})'.format(loss, acc, running_corrects, example_count))\n",
    "    return loss, acc\n",
    "\n",
    "\n",
    "def train_model(model, data_loaders, criterion, optimizer, scheduler, save_dir, num_epochs=25, use_gpu=False):\n",
    "    print('Training Model with use_gpu={}...'.format(use_gpu))\n",
    "    since = time.time()\n",
    "\n",
    "    best_model_wts = model.state_dict()\n",
    "    best_acc = 0.0\n",
    "    for epoch in range(num_epochs):\n",
    "        print('Epoch {}/{}'.format(epoch, num_epochs - 1))\n",
    "        print('-' * 10)\n",
    "        train_begin = time.time()\n",
    "        train_loss, train_acc = train(model, data_loaders['train'], criterion, optimizer, use_gpu)\n",
    "        train_time = time.time() - train_begin\n",
    "        print('Epoch Train Time: {:.0f}m {:.0f}s'.format(train_time // 60, train_time % 60))\n",
    "        \n",
    "        validation_begin = time.time()\n",
    "        val_loss, val_acc = validate(model, data_loaders['validation'], criterion, use_gpu)\n",
    "        validation_time = time.time() - validation_begin\n",
    "        print('Epoch Validation Time: {:.0f}m {:.0f}s'.format(validation_time // 60, validation_time % 60))\n",
    "        \n",
    "        # deep copy the model\n",
    "        is_best = val_acc > best_acc\n",
    "        if is_best:\n",
    "            best_acc = val_acc\n",
    "            best_model_wts = model.state_dict()\n",
    "\n",
    "        save_checkpoint(save_dir, {\n",
    "            'epoch': epoch,\n",
    "            'best_acc': best_acc,\n",
    "            'state_dict': model.state_dict(),\n",
    "            # 'optimizer': optimizer.state_dict(),\n",
    "        }, is_best)\n",
    "\n",
    "        scheduler.step()\n",
    "\n",
    "    time_elapsed = time.time() - since\n",
    "    print('Training complete in {:.0f}m {:.0f}s'.format(\n",
    "        time_elapsed // 60, time_elapsed % 60))\n",
    "    print('Best val Acc: {:4f}'.format(best_acc))\n",
    "    # load best model weights\n",
    "    model.load_state_dict(best_model_wts)\n",
    "\n",
    "    return model\n",
    "\n",
    "\n",
    "def save_checkpoint(save_dir, state, is_best):\n",
    "    savepath = save_dir + '/' + 'checkpoint.pth.tar'\n",
    "    torch.save(state, savepath)\n",
    "    if is_best:\n",
    "        shutil.copyfile(savepath, save_dir + '/' + 'model_best.pth.tar')\n",
    "\n",
    "\n",
    "def test_model(model, test_loader, use_gpu=False):\n",
    "    model.eval()  # Set model to evaluate mode\n",
    "    running_corrects = 0\n",
    "    example_count = 0\n",
    "    test_begin = time.time()\n",
    "    # Iterate over data.\n",
    "    with torch.no_grad():\n",
    "        for batch in test_loader:\n",
    "            sentences = batch.text[0].transpose(1, 0)\n",
    "            tags = batch.tags.transpose(1, 0)\n",
    "            ''' Similar to dev, except do we not need to compute the loss here\n",
    "            '''\n",
    "            # TODO\n",
    "            \n",
    "\n",
    "    acc = (running_corrects / example_count) * 100\n",
    "    print('Test Acc: {:2.3f} ({}/{})'.format(acc, running_corrects, example_count))\n",
    "    test_time = time.time() - test_begin\n",
    "    print('Test Time: {:.0f}m {:.0f}s'.format(test_time // 60, test_time % 60))\n",
    "    return acc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class POSTagger(nn.Module):\n",
    "    def __init__(self, rnn_class, embedding_dim, hidden_dim, vocab_size, target_size, num_layers):\n",
    "        super(POSTagger, self).__init__()\n",
    "        ''' Define your model here\n",
    "            Basically, your model only need three components:\n",
    "            - an embedding layer\n",
    "            - a bidirectional RNN (LSTM, GRU) that takes the embeddings and outputs hidden states\n",
    "            - a final linear prediction layer to convert hidden states to tag scores\n",
    "            Optionally, define extra layers such as dropout to prevent overfitting. \n",
    "        ''' \n",
    "        # TODO\n",
    "        \n",
    "        \n",
    "    def forward(self, sentences):\n",
    "        ''' Define your forward method\n",
    "        ''' \n",
    "        # TODO\n",
    "        \n",
    "        \n",
    "        return tag_scores\n",
    "    \n",
    "model = POSTagger(\"lstm\", EMBEDDING_DIM, HIDDEN_DIM, text_vocab_size, tag_vocab_size, 3)\n",
    "if USE_CUDA:\n",
    "    model = model.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "LR = 0.001\n",
    "GAMMA = 1.\n",
    "STEP_SIZE = 10\n",
    "NUM_EPOCHS = 10\n",
    "SAVE_DIR = \"./save/\"\n",
    "loss_fn = nn.CrossEntropyLoss(size_average=False)\n",
    "optimizer = optim.Adam(model.parameters(), lr=LR)\n",
    "exp_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=STEP_SIZE, gamma=GAMMA)\n",
    "model = train_model(model, dataloaders, loss_fn, optimizer, exp_lr_scheduler, SAVE_DIR, NUM_EPOCHS, use_gpu=USE_CUDA)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
